# Code copied from pytorch-tutorial https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/04-utils/tensorboard/logger.py 
# import tensorflow as tf
from tensorboardX import SummaryWriter
import numpy as np


class TensorBoard(object):
    
    def __init__(self, log_dir):
        """Create a summary writer logging to log_dir."""
        self.writer = SummaryWriter(log_dir)

    def scalar_summary(self, tag, value, step):
        """Log a scalar variable."""
        self.writer.add_scalar(tag, value, global_step=step)

    def image_summary(self, tag, images, step):
        """Log a list of images."""
        for img in images:
            assert type(img) is np.ndarray
            img = np.transpose(img, (2,0,1))    # 从HWC转换成CHW格式
            self.writer.add_image(tag, img, global_step=step)

    def histo_summary(self, tag, values, step, bins=1000):
        """Log a histogram of the tensor of values."""
        self.writer.add_histogram(tag, values, global_step=step, bins=bins)
